#!/usr/bin/env python3

# Copyright (c) 2005-2014 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import path
import copy
import kmod
import glob
import pathlib
import robust_layer.simple_fops
from ._util import Util
from ._po import Arch
from ._po import KernelType
from ._exception import BootEntryInvalid


class BootEntrySpec:

    def __init__(self, dict_data):
        self._assertDictData(dict_data)
        self._data = dict_data

    @property
    def arch(self):
        return Arch(self._data["arch"])

    @property
    def kernel_type(self):
        return KernelType(self._data["kernel_type"])

    @property
    def kernel_addons(self):
        return self._data["kernel_addon_params"].keys()

    def get_kernel_params(self):
        return self._data["kernel_params"]

    def get_kernel_addon_params(self, kernel_addon_name):
        return self._data["kernel_addon_params"][kernel_addon_name]

    def get_initramfs_params(self):
        return self._data["initramfs_params"]

    def to_dict_data(self):
        return copy.deepcopy(self._data)

    def _assertDictData(self, dict_data):
        if "arch" not in dict_data:
            assert False
        if dict_data["arch"] not in [x.value for x in list(Arch)]:
            assert False

        if "kernel_type" not in dict_data:
            assert False
        if dict_data["kernel_type"] not in [x.value for x in list(KernelType)]:
            assert False

        if "kernel_params" not in dict_data:
            assert False
        if not isinstance(dict_data["kernel_params"], dict):
            assert False

        if "kernel_addon_params" not in dict_data:
            assert False
        for kernel_addon_name, params in dict_data["kernel_addon_params"].items():
            if not isinstance(params, dict):
                assert False

        if "initramfs_params" not in dict_data:
            assert False
        if not isinstance(dict_data["initramfs_params"], dict):
            assert False


class BootEntry:

    def __init__(self, bbki, arch, verstr, history_entry=False):
        self._bbki = bbki
        self._verstr = verstr

        if not history_entry:
            self._bootDir = self._bbki._fsLayout.get_boot_dir()
        else:
            self._bootDir = self._bbki._fsLayout.get_boot_history_dir()

        # get spec file and postfix
        postfix = arch + "-" + self._verstr
        specFile = os.path.join(self._bootDir, "spec-" + postfix)
        if not os.path.exists(specFile):
            raise BootEntryInvalid("no boot entry spec file found in %s" % (self._bootDir))

        # postfix example: x86_64-3.9.11-gentoo-r1
        partList = postfix.split("-")
        if len(partList) < 2:
            raise BootEntryInvalid("illegal postfix found in %s" % (path))
        if not Util.isValidKernelArch(partList[0]):                             # FIXME: isValidKernelArch should be moved out from util
            raise BootEntryInvalid("illegal postfix found in %s" % (path))
        if not Util.isValidKernelVer(partList[1]):                              # FIXME: isValidKernelVer should be moved out from util
            raise BootEntryInvalid("illegal postfix found in %s" % (path))

        # parse spec file
        self._spec = None
        self._kernelVer = None
        self._addonVerDict = None
        if True:
            kernelParams = None
            addonParamsDict = None
            initramfsParams = None
            kernelVer = None
            addonVerDict = None

            lineList = pathlib.Path(specFile).read_text().split("\n")

            def __parseParameters(myName, s):
                ret = {}
                if s != "":
                    for item in s.split(" "):
                        tlist = item.split("=")
                        if len(tlist) != 2:
                            raise BootEntryInvalid("invalid %s parameter" % (myName))
                        ret[tlist[0]] = tlist[1]
                return ret

            for line in lineList:
                m = re.fullmatch("kernel-params: (.*)", line)
                if m is not None:
                    if kernelParams is not None:
                        raise BootEntryInvalid("duplicate \"kernel-params\" found in %s" % (specFile))
                    kernelParams = __parseParameters("kernel", m.group(1))

            addonParamsDict = {}
            for line in lineList:
                m = re.fullmatch("kernel-addon-(\\S+)-params: (.*)", line)
                if m is not None:
                    if m.group(1) in addonParamsDict:
                        raise BootEntryInvalid("duplicate \"kernel-addons-%s-params\" found in %s" % (m.group(1), specFile))
                    addonParamsDict[m.group(1)] = __parseParameters("kernel addon", m.group(2))

            for line in lineList:
                m = re.fullmatch("initramfs-params: (.*)", line)
                if m is not None:
                    if initramfsParams is not None:
                        raise BootEntryInvalid("duplicate \"initramfs-params\" found in %s" % (specFile))
                    initramfsParams = __parseParameters("initramfs", m.group(1))

            for line in lineList:
                m = re.fullmatch("kernel-version: (.*)", line)
                if m is not None:
                    if kernelVer is not None:
                        raise BootEntryInvalid("duplicate \"kernel-version\" found in %s" % (specFile))
                    kernelVer = m.group(1)

            addonVerDict = {}
            for line in lineList:
                m = re.fullmatch("kernel-addon-(\\S+)-version: (.*)", line)
                if m is not None:
                    if m.group(1) in addonVerDict:
                        raise BootEntryInvalid("duplicate \"kernel-addon-%s-version\" found in %s" % (m.group(1), specFile))
                    addonVerDict[m.group(1)] = m.group(2)

            if kernelParams is None:
                raise BootEntryInvalid("no \"kernel-params\" found in %s" % (specFile))

            if addonParamsDict.keys() != addonVerDict.keys():
                raise BootEntryInvalid("inconsistent \"kernel-addons-params\" and \"kernel-addons-version\" found in %s" % (specFile))

            if initramfsParams is None:
                raise BootEntryInvalid("no \"initramfs-params\" found in %s" % (specFile))

            if kernelVer is None:
                raise BootEntryInvalid("no \"kernel-version\" found in %s" % (specFile))

            if verstr != kernelVer and verstr.startswith(kernelVer + "-"):
                raise BootEntryInvalid("invalid \"kernel-version\" in %s" % (specFile))

            self._spec = BootEntrySpec({
                "arch": arch,
                "kernel_type": KernelType.LINUX.value,
                "kernel_params": kernelParams,
                "kernel_addon_params": addonParamsDict,
                "initramfs_params": initramfsParams,
            })
            self._kernelVer = kernelVer
            self._addonVerDict = addonVerDict

    @property
    def spec(self):
        return self._spec

    @property
    def postfix(self):
        # string, eg: "x86_64-3.9.11-gentoo-r1"
        return self._spec.arch.value + "-" + self._verstr

    @property
    def verstr(self):
        # string, eg: "3.9.11-gentoo-r1"
        return self._verstr

    @property
    def spec_filename(self):
        # string, eg: "spec-x86_64-3.9.11-gentoo-r1"
        return "spec-" + self.postfix

    @property
    def spec_filepath(self):
        # string, eg: "/boot/spec-x86_64-3.9.11-gentoo-r1"
        return os.path.join(self._bootDir, self.spec_filename)

    @property
    def kernel_filename(self):
        # string, eg: "kernel-x86_64-3.9.11-gentoo-r1"
        return "kernel-" + self.postfix

    @property
    def kernel_filepath(self):
        # string, eg: "/boot/kernel-x86_64-3.9.11-gentoo-r1"
        return os.path.join(self._bootDir, self.kernel_filename)

    @property
    def kernel_config_filename(self):
        # string, eg: "config-x86_64-3.9.11-gentoo-r1"
        return "config-" + self.postfix

    @property
    def kernel_config_filepath(self):
        # string, eg: "/boot/config-x86_64-3.9.11-gentoo-r1"
        return os.path.join(self._bootDir, self.kernel_config_filename)

    @property
    def kernel_modules_dirpath(self):
        # string, eg: "/lib/modules/5.1.14-gentoo-r1"
        return self._bbki._fsLayout.get_kernel_modules_dir(self._verstr)

    @property
    def firmware_dirpath(self):
        # string, eg: "/lib/firmware"
        return self._bbki._fsLayout.get_firmware_dir()

    @property
    def initramfs_filename(self):
        # string, eg: "initramfs-x86_64-3.9.11-gentoo-r1"
        return "initramfs-" + self.postfix

    @property
    def initramfs_filepath(self):
        # string, eg: "/boot/initramfs-x86_64-3.9.11-gentoo-r1"
        return os.path.join(self._bootDir, self.initramfs_filename)

    @property
    def initramfs_tar_filename(self):
        # string, eg: "initramfs-files-x86_64-3.9.11-gentoo-r1.tar.bz2"
        return "initramfs-files-" + self.postfix + ".tar.bz2"

    @property
    def initramfs_tar_filepath(self):
        # string, eg: "/boot/initramfs-x86_64-3.9.11-gentoo-r1.tar.bz2"
        return os.path.join(self._bootDir, self.initramfs_tar_filename)

    def is_historical(self):
        return (self._bootDir == self._bbki._fsLayout.get_boot_history_dir())

    def get_kernel_version(self):
        # string, eg: "3.9.11"
        return self._kernelVer

    def get_kernel_addon_version(self, kernel_addon_name):
        return self._addonVerDict[kernel_addon_name]

    def has_kernel_files(self):
        if not os.path.exists(self.spec_filepath):
            return False
        if not os.path.exists(self.kernel_filepath):
            return False
        if not os.path.exists(self.kernel_config_filepath):
            return False
        if not os.path.exists(self.kernel_modules_dirpath):
            return False
        if not os.path.exists(self.firmware_dirpath):
            return False
        return True

    def has_initramfs_files(self):
        if not os.path.exists(self.initramfs_filepath):
            return False
        if not os.path.exists(self.initramfs_tar_filepath):
            return False
        return True

    def __eq__(self, other):
        if type(self) is not type(other):
            return False
        if self._spec.arch != other._spec.arch:
            return False
        if self._verstr != other._verstr:
            return False
        if self._bootDir != other._bootDir:
            return False
        return True

    def __hash__(self):
        return hash((self._spec.arch, self._verstr, self._bootDir))


class BootEntryUtils:

    def __init__(self, bbki):
        self._bbki = bbki

    @property
    def src_arch(self, bootEntry):
        # FIXME: what's the difference with arch?

        if bootEntry.spec.arch.value == "i386" or bootEntry.spec.arch.value == "x86_64":
            return "x86"
        elif bootEntry.spec.arch.value == "sparc32" or bootEntry.spec.arch.value == "sparc64":
            return "sparc"
        elif bootEntry.spec.arch.value == "sh":
            return "sh64"
        else:
            return bootEntry.spec.arch.value

    def get_filepaths(self, bootEntry):
        return [
            bootEntry.spec_filepath,
            bootEntry.kernel_filepath,
            bootEntry.kernel_config_filepath,
            bootEntry.initramfs_filepath,
            bootEntry.initramfs_tar_filepath,
        ]

    def get_kmod_filenames_by_alias(self, bootEntry, kmod_alias, with_deps=False):
        return [x[len(bootEntry.kernel_modules_dirpath):] for x in self.get_kmod_filepaths_by_alias(bootEntry, kmod_alias, with_deps)]

    def get_kmod_filepaths_by_alias(self, bootEntry, kmod_alias, with_deps=False):
        kmodList = dict()                                                           # use dict to remove duplication while keeping order
        ctx = kmod.Kmod(bootEntry.kernel_modules_dirpath.encode("utf-8"))           # FIXME: why encode is neccessary?
        self._getKmodAndDeps(ctx, kmod_alias, with_deps, kmodList)
        return list(kmodList)

    def get_firmware_filenames_by_kmod(self, bootEntry, kmod_filepath):
        # python-kmod bug: can only recognize the last firmware in modinfo
        # so use the command output of modinfo directly
        ret = []
        for line in Util.cmdCall("/bin/modinfo", kmod_filepath).split("\n"):
            m = re.fullmatch("firmware: +(\\S.*)", line)
            if m is not None:
                ret.append(m.group(1))
        return ret

    def get_firmware_filepaths_by_kmod(self, bootEntry, kmod_filepath):
        return [os.path.join(bootEntry.firmware_dirpath, x) for x in self.get_firmware_filenames_by_kmod(bootEntry, kmod_filepath)]

    def get_firmware_filenames(self, bootEntry):
        ret = set()

        for fullKoFn in glob.glob(os.path.join(bootEntry.kernel_modules_dirpath, "**", "*.ko"), recursive=True):
            ret |= set(self.get_firmware_filenames_by_kmod(bootEntry, fullKoFn))

        fwExtFileRecordFn = os.path.join(bootEntry.kernel_modules_dirpath, "firmware.extra-files")
        if os.path.exists(fwExtFileRecordFn):
            ret |= set(Util.readListFile(fwExtFileRecordFn))

        return sorted(list(ret))

    def get_firmware_filepaths(self, bootEntry):
        return [os.path.join(bootEntry.firmware_dirpath, x) for x in self.get_firmware_filenames(bootEntry)]

    def move_to_history(self, bootEntry):
        robust_layer.simple_fops.mkdir(self._bbki._fsLayout.get_boot_history_dir())
        for fullfn in self.get_filepaths(bootEntry):
            if os.path.exists(fullfn):
                robust_layer.simple_fops.mv_to_dir(fullfn, self._bbki._fsLayout.get_boot_history_dir())

    def new_from_postfix(self, postfix, history_entry=False):
        # postfix example: x86_64-3.9.11-gentoo-r1
        partList = postfix.split("-")
        if len(partList) < 2:
            raise ValueError("illegal postfix")
        if not Util.isValidKernelArch(partList[0]):         # FIXME: isValidKernelArch should be moved out from util
            raise ValueError("illegal postfix")
        if not Util.isValidKernelVer(partList[1]):          # FIXME: isValidKernelVer should be moved out from util
            raise ValueError("illegal postfix")

        arch = partList[0]
        verstr = "-".join(partList[1:])
        return BootEntry(self._bbki, arch, verstr, history_entry)

    def getRedundantKernelModulesDirs(self, bootEntryList):
        if not os.path.exists(self._bbki._fsLayout.get_kernel_modules_dir()):
            return []

        ret = glob.glob(os.path.join(self._bbki._fsLayout.get_kernel_modules_dir(), "*"))  # mark /lib/modules/* (no recursion) as to-be-deleted
        for be in bootEntryList:
            try:
                ret.remove(self._bbki._fsLayout.get_kernel_modules_dir(be.verstr))         # don't delete files belongs to a boot-entry
            except ValueError:
                pass
        return ret

    def getRedundantFirmwareFiles(self, bootEntryList):
        if not os.path.exists(self._bbki._fsLayout.get_firmware_dir()):
            return []

        tset = set(Util.globDirRecursively(self._bbki._fsLayout.get_firmware_dir()))   # mark /lib/firmware/* (recursive) as to-be-deleted
        for be in bootEntryList:
            for fp in BootEntryUtils(self._bbki).get_firmware_filepaths(be):
                while fp != "/":
                    tset.discard(fp)
                    fp = os.path.dirname(fp)                                           # don't delete files (and it ancestor directories) belongs to a boot-entry
        return sorted(list(tset))

    def _getKmodAndDeps(self, ctx, kmodAlias, withDeps, result):
        kmodObjList = list(ctx.lookup(kmodAlias))
        if len(kmodObjList) > 0:
            assert len(kmodObjList) == 1
            kmodObj = kmodObjList[0]

            if withDeps and "depends" in kmodObj.info and kmodObj.info["depends"] != "":
                for kmodAlias in kmodObj.info["depends"].split(","):
                    self._getKmodAndDeps(ctx, kmodAlias, withDeps, result)

            if kmodObj.path is not None:
                # this module is not built into the kernel
                result[kmodObj.path] = None
